import cv2
import numpy as np
import sys


def main():
    # 检查命令行参数
    if len(sys.argv) != 2:
        print("使用方法: python red_armor_detector.py <图片路径>")
        sys.exit(1)

    # 获取图片路径
    image_path = sys.argv[1]

    # 加载图片
    image = cv2.imread(image_path)
    if image is None:
        print(f"错误: 无法加载图片 '{image_path}'")
        sys.exit(1)

    # 将图片从BGR转换到HSV色彩空间
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    # 设定红色的HSV阈值范围
    lower_red1 = np.array([0, 100, 100])
    upper_red1 = np.array([10, 255, 255])
    lower_red2 = np.array([160, 100, 100])
    upper_red2 = np.array([180, 255, 255])

    # 创建掩码
    mask1 = cv2.inRange(hsv_image, lower_red1, upper_red1)
    mask2 = cv2.inRange(hsv_image, lower_red2, upper_red2)
    red_mask = cv2.bitwise_or(mask1, mask2)

    # 显示原始图片和处理后的掩码
    cv2.imshow('Original Image', image)
    cv2.imshow('Red Mask', red_mask)

    print("按任意键关闭窗口...")
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()